'''
author:        Wang Chenyang <cy-wang21@mails.tsinghua.edu.cn>
date:          2025-08-13 18:46:20
Copyright © Department of Physics, Tsinghua University. All rights reserved
'''

import matplotlib.pyplot as plt
import numpy as np

''' 备注：
    纸条与床头柜接缝长度：102.0 mm, 像素值82.05485969764375
    纸条露出部分上边缘长度：146.7 mm, 像素值112.75194011634566
    纸条露出部分下边缘，最外层纸片长度：141.4 mm, 像素值106.38138934982942
    比值：1.243070798924649, 1.3010862593461738,  1.329179858095421
    平均比值：1.2911123054554148
'''

MM_PER_PX = 1.243070798924649


def plot_x_y_curves(data_file):
    data = np.loadtxt(data_file, delimiter=',', skiprows=1)
    time = data[:, 0]
    x = data[:, 1]
    y = data[:, 2]

    k_x_tsq = np.polyfit(time ** 2, x, 1)
    k_y_tsq = np.polyfit(time ** 2, y, 1)

    plt.figure()
    plt.plot(time, x, 'b.')
    plt.plot(time, y, 'r.')
    plt.legend(['x', 'y'])
    plt.show()


def 探究落体运动规律(data_file):
    data = np.loadtxt(data_file, delimiter=',', skiprows=1)
    time = data[:, 0]
    x = data[:, 1]
    y = data[:, 2]

    k_x_tsq = np.polyfit(time ** 2, x, 1)
    k_y_tsq = np.polyfit(time ** 2, y, 1)

    plt.figure()
    plt.plot(time ** 2, x, 'b.')
    plt.plot(time ** 2, k_x_tsq[0] * time ** 2 + k_x_tsq[1], 'b-')
    plt.plot(time ** 2, y, 'r.')
    plt.plot(time ** 2, k_y_tsq[0] * time ** 2 + k_y_tsq[1], 'r-')
    plt.legend(['x', 'x fit', 'y', 'y fit'])
    print(k_x_tsq)
    print(k_y_tsq)
    a_x = k_x_tsq[0] * 2 * MM_PER_PX / 1000
    a_y = k_y_tsq[0] * 2 * MM_PER_PX / 1000
    print("x的加速度为:", a_x, "m/s^2")
    print("y的加速度为:", a_y, "m/s^2")
    print("加速度矢量为:", np.sqrt(a_x ** 2 + a_y ** 2), "m/s^2")
    plt.show()


def 拟合重力加速度(data_file):
    data = np.loadtxt(data_file, delimiter=',', skiprows=1)
    time = data[:, 0]
    x = data[:, 1]
    y = data[:, 2]

    x_over_t = np.polyfit(time, x, 2)
    y_over_t = np.polyfit(time, y, 2)
    a_x = x_over_t[0] * 2 * MM_PER_PX / 1000
    a_y = y_over_t[0] * 2 * MM_PER_PX / 1000
    print("x的加速度为:", a_x, "m/s^2")
    print("y的加速度为:", a_y, "m/s^2")
    print("加速度矢量为:", np.sqrt(a_x ** 2 + a_y ** 2), "m/s^2")

    plt.plot(time, x, 'b.')
    plt.plot(time, x_over_t[0] * time ** 2 + x_over_t[1] * time + x_over_t[2], 'b-')
    plt.plot(time, y, 'r.')
    plt.plot(time, y_over_t[0] * time ** 2 + y_over_t[1] * time + y_over_t[2], 'r-')
    plt.legend(['x', 'x fit', 'y', 'y fit'])
    plt.show()


if __name__ == "__main__":
    # plot_x_y_curves(data_file)
    for j in range(5):
        data_file = f"data/free_fall_data_{j}.csv"
        拟合重力加速度(data_file)


